package com.github.elliot.gatewaycenter.domain;

import lombok.extern.slf4j.Slf4j;
import org.reactivestreams.Publisher;
import org.springframework.core.io.buffer.DataBuffer;
import org.springframework.core.io.buffer.DataBufferUtils;
import org.springframework.http.server.reactive.ServerHttpResponse;
import org.springframework.http.server.reactive.ServerHttpResponseDecorator;
import reactor.core.publisher.Flux;
import reactor.core.publisher.Mono;

import javax.crypto.Cipher;
import javax.crypto.spec.SecretKeySpec;
import java.nio.charset.StandardCharsets;
import java.util.Base64;

@Slf4j
public class EncryptedServerHttpResponse extends ServerHttpResponseDecorator {

    public EncryptedServerHttpResponse(ServerHttpResponse delegate) {
        super(delegate);
    }

    @Override
    public Mono<Void> writeWith(Publisher<? extends DataBuffer> body) {
        // 1. 收集原始响应体并转换为字符串
        Mono<String> originalBody = DataBufferUtils.join(body)
                .map(dataBuffer -> {
                    byte[] bytes = new byte[dataBuffer.readableByteCount()];
                    dataBuffer.read(bytes);
                    DataBufferUtils.release(dataBuffer); // 释放缓冲区
                    return new String(bytes, StandardCharsets.UTF_8);
                });

        // 2. 加密响应体字符串
        Mono<DataBuffer> encryptedBody = originalBody
                .map(originalText -> {
                    try {
                        // 使用 AES 加密
                        String encryptedText = encryptAES(originalText, "aaa");
                        log.info(encryptedText);
                        // 将加密后的字符串包装为 DataBuffer
                        byte[] bytes = encryptedText.getBytes(StandardCharsets.UTF_8);
                        return getDelegate().bufferFactory().wrap(bytes);
                    } catch (Exception e) {
                        throw new RuntimeException("Response encryption failed", e);
                    }
                });

        // 3. 将加密后的 DataBuffer 写入响应
        return super.writeWith(encryptedBody.flux());
    }

    private String encryptAES(String data, String key) throws Exception {
        // 创建 AES 密钥规范
        SecretKeySpec keySpec = new SecretKeySpec(key.getBytes(StandardCharsets.UTF_8), "AES");
        // 获取 AES Cipher 实例（这里使用 ECB 模式，PKCS5Padding 填充）
        Cipher cipher = Cipher.getInstance("AES/ECB/PKCS5Padding");
        cipher.init(Cipher.ENCRYPT_MODE, keySpec);
        // 执行加密并返回 Base64 编码的字符串
        byte[] encryptedBytes = cipher.doFinal(data.getBytes(StandardCharsets.UTF_8));
        return Base64.getEncoder().encodeToString(encryptedBytes);
    }
}